#include "x64Emitter.h"


x64Emitter::x64Emitter(void* pvBuffer, int nSize)
{
    m_pbBuffer = static_cast<uint8_t*>(pvBuffer);
    m_nBufferSize = nSize;

    m_pbWrite = m_pbBuffer;
}

x64Emitter::~x64Emitter()
{
}

// ----------------------------------------------

GPRegister::~GPRegister()
{
}

GPRegister GPRegister::operator+(const GPRegister& right) const
{
    assert(!mul || !right.mul);
    
    GPRegister result(-1);

    if (!mul)
        result = GPRegister(id, right.id, right.mul, immOffset + right.immOffset);
    else
        result = GPRegister(right.id, id, mul, immOffset + right.immOffset);

    return result;
}

GPRegister GPRegister::operator+(const int& right) const
{
    GPRegister result(id, idOffset, mul, right + immOffset);

    return result;
}

GPRegister GPRegister::operator-(const int& right) const
{
    GPRegister result(id, idOffset, mul, immOffset - right);

    return result;
}

GPRegister GPRegister::operator*(const int& nMul) const
{
    assert(idOffset < 0);

    GPRegister result(id, idOffset, nMul, immOffset * nMul);

    return result;
}

MemoryAddress GPRegister::operator*() const
{
    MemoryAddress result(id, idOffset, mul, immOffset);

    return result;
}

// ----------------------------------------------

template<typename opcode_t>
void x64Emitter::instructionRegToReg(const GPRegister& rd, const GPRegister& rs, opcode_t opcode)
{
    const uint8_t RMMemExtend = rd > RDI ? 0x01 : 0x00;
    const uint8_t RMRegExtend = rs > RDI ? 0x04 : 0x00;

    const uint8_t RMMem = rd.id & 0x07;
    const uint8_t RMReg = rs.id & 0x07;

    *m_pbWrite++ = 0x48 + RMMemExtend + RMRegExtend;
    *(reinterpret_cast<opcode_t*>(m_pbWrite)) = opcode;
    m_pbWrite += sizeof(opcode_t);
    *m_pbWrite++ = 0xC0 + (RMReg << 3) + RMMem;
}

// ----------------------------------------------

template<typename opcode_t>
void x64Emitter::instructionMemToReg(const GPRegister& rd, const MemoryAddress& mem, opcode_t opcode)
{
    const uint8_t RMRegExtend = rd > RDI ? 0x04 : 0x00;
    const uint8_t RMMemExtend = mem.id > RDI.id ? 0x01 : 0x00;
    const uint8_t RMMemExtendOff = mem.idOffset > RDI.id ? 0x02 : 0x00;

    const bool isSIB = (mem.idOffset >= 0) || ((mem.id & 0x07) == RSP.id);

    const uint8_t RMReg = rd.id & 0x07;
    const uint8_t RMMem = isSIB ? 0x04 : mem.id & 0x07;

    
    const bool is32BitsDisp = (mem.immOffset > 127) || (mem.immOffset < -128);

    uint8_t RMDisp = mem.immOffset ? (is32BitsDisp ? 0x80 : 0x40) : 0x00;

    // Special case, you can't code [rbp], so we code [rbp + 0] like GCC does
    const bool isRbpCase = (mem.id == RBP.id && mem.idOffset < 0 && mem.immOffset == 0);

    if (isRbpCase)
        RMDisp = 0x40;

    *m_pbWrite++ = 0x48 + RMMemExtend + RMRegExtend + RMMemExtendOff;    // REX
    *(reinterpret_cast<opcode_t*>(m_pbWrite)) = opcode;                  
    m_pbWrite += sizeof(opcode_t);
    *m_pbWrite++ = RMDisp + (RMReg << 3) + RMMem;       // ModR/M

    if (isSIB)
    {
        const uint8_t base = mem.id & 0x07;
        const uint8_t offset = mem.idOffset >= 0 ? mem.idOffset  & 0x07 : 0x04;
        uint8_t SIBMul;

        switch (mem.mul)
        {
        case 0:
        case 1:
            SIBMul = 0x00; break;
        case 2:
            SIBMul = 0x40; break;
        case 4:
            SIBMul = 0x80; break;
        case 8:
            SIBMul = 0xC0; break;
        }

        *m_pbWrite++ = SIBMul + (offset << 3) + base;   // SIB
    }

    if (mem.immOffset || isRbpCase)
    {
        if (is32BitsDisp)
        {
            *m_pbWrite++ = mem.immOffset & 0xFF;
            *m_pbWrite++ = (mem.immOffset >> 8) & 0xFF;
            *m_pbWrite++ = (mem.immOffset >> 16) & 0xFF;
            *m_pbWrite++ = (mem.immOffset >> 24) & 0xFF;
        }
        else
        {
            *m_pbWrite++ = static_cast<uint8_t>(mem.immOffset);
        }
    }
}

// ----------------------------------------------

#define GPR_MEMORY_INSTRUCTION(name, opcodeRR, opcodeRM) \
    void x64Emitter::name(const GPRegister& rd, const GPRegister& rs) { instructionRegToReg(rd, rs, (uint8_t)opcodeRR); } \
    void x64Emitter::name(const MemoryAddress& mem, const GPRegister& rs) { instructionMemToReg(rs, mem, (uint8_t)opcodeRR); } \
    void x64Emitter::name(const GPRegister& rd, const MemoryAddress& mem) { instructionMemToReg(rd, mem, (uint8_t)opcodeRM); }


GPR_MEMORY_INSTRUCTION(MOV, 0x89, 0x8B)
GPR_MEMORY_INSTRUCTION(ADD, 0x01, 0x03)
GPR_MEMORY_INSTRUCTION(ADC, 0x11, 0x13)
GPR_MEMORY_INSTRUCTION(AND, 0x21, 0x23)
GPR_MEMORY_INSTRUCTION(XOR, 0x31, 0x33)
GPR_MEMORY_INSTRUCTION(OR , 0x09, 0x0B)
GPR_MEMORY_INSTRUCTION(SBB, 0x19, 0x1B)
GPR_MEMORY_INSTRUCTION(SUB, 0x29, 0x2B)
GPR_MEMORY_INSTRUCTION(CMP, 0x39, 0x3B)
GPR_MEMORY_INSTRUCTION(TEST, 0x85, 0x85)
GPR_MEMORY_INSTRUCTION(XCHG, 0x87, 0x87)

#undef GPR_MEMORY_INSTRUCTION

void x64Emitter::LEA (const GPRegister& rd, const MemoryAddress& mem) { instructionMemToReg(rd, mem, (uint8_t)0x8D); }
void x64Emitter::IMUL(const GPRegister& rd, const MemoryAddress& mem) { instructionMemToReg(rd, mem, (uint16_t)0xAF0F); }
void x64Emitter::IMUL(const GPRegister& rd, const GPRegister& rs) { instructionRegToReg(rd, rs, (uint16_t)0xAF0F); }
void x64Emitter::IMUL(const GPRegister& rd, const MemoryAddress& mem, int32_t imm32) 
{
    instructionMemToReg(rd, mem, (uint8_t)0x69);

    *(reinterpret_cast<int32_t*>(m_pbWrite)) = imm32;
    m_pbWrite += 4;
}
void x64Emitter::IMUL(const GPRegister& rd, const GPRegister& rs, int32_t imm32) 
{
    instructionRegToReg(rd, rs, (uint8_t)0x69);

    *(reinterpret_cast<int32_t*>(m_pbWrite)) = imm32;
    m_pbWrite += 4;
}

// ----------------------------------------------

void x64Emitter::NOT(const GPRegister& rd) { group3(rd, 0x02); }
void x64Emitter::NEG(const GPRegister& rd) { group3(rd, 0x03); }

// ----------------------------------------------

// MOV Rx, imm64
// Also known as MOVABS
// No RM byte, Reg coded in the opcode

void x64Emitter::MOV(const GPRegister& rd, const int64_t& imm64)
{
    const uint8_t RMRegExtend = rd > RDI ? 0x01 : 0x00;
    const uint8_t RMReg = rd.id & 0x07;

    *m_pbWrite++ = 0x48 + RMRegExtend;
    *m_pbWrite++ = 0xB8 + RMReg;
    
    for (int i = 0; i < 8; i++)
        *m_pbWrite++ = (imm64 >> (i * 8)) & 0xFF;
}

// ----------------------------------------------

void x64Emitter::MOV(const GPRegister& rd, const int32_t& imm32)
{
    const uint8_t RMRegExtend = rd > RDI ? 0x01 : 0x00;
    const uint8_t RMReg = rd.id & 0x07;

    const bool is32BitsImm = (imm32 > 127) || (imm32 < -128);

    *m_pbWrite++ = 0x48 + RMRegExtend;
    *m_pbWrite++ = is32BitsImm ? 0xC7 : 0xC6;
    *m_pbWrite++ = 0xC0 + RMReg;
    
    if (is32BitsImm)
    {
        for (int i = 0; i < 4; i++)
            *m_pbWrite++ = (imm32 >> (i * 8)) & 0xFF;
    }
    else
    {
        *m_pbWrite++ = imm32 & 0xFF;
    }
}

// ----------------------------------------------

// Immediate Group 1

void x64Emitter::immediateGroup1(const GPRegister& rd, const int32_t& imm32, uint8_t opcodeExtend)
{
    const uint8_t RMRegExtend = rd > RDI ? 0x01 : 0x00;
    const uint8_t RMReg = rd.id & 0x07;

    const bool is32BitsImm = (imm32 > 127) || (imm32 < -128);

    *m_pbWrite++ = 0x48 + RMRegExtend + (opcodeExtend << 3);
    *m_pbWrite++ = is32BitsImm ? 0x81 : 0x83;
    *m_pbWrite++ = 0xC0 + RMReg;

    if (is32BitsImm)
    {
        for (int i = 0; i < 4; i++)
            *m_pbWrite++ = (imm32 >> (i * 8)) & 0xFF;
    }
    else
    {
        *m_pbWrite++ = imm32 & 0xFF;
    }
}

// ----------------------------------------------

// 1 Operand - Group 3

void x64Emitter::group3(const GPRegister& rd, uint8_t opcodeExtend)
{
    const uint8_t RMRegExtend = rd > RDI ? 0x01 : 0x00;
    const uint8_t RMReg = rd.id & 0x07;

    *m_pbWrite++ = 0x48 + RMRegExtend;
    *m_pbWrite++ = 0xF7;
    *m_pbWrite++ = 0xC0 + RMReg + (opcodeExtend << 3);
}

// ----------------------------------------------

void x64Emitter::ADD(const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x00); }
void x64Emitter::OR (const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x01); }
void x64Emitter::ADC(const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x02); }
void x64Emitter::SBB(const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x03); }
void x64Emitter::AND(const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x04); }
void x64Emitter::SUB(const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x05); }
void x64Emitter::XOR(const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x06); }
void x64Emitter::CMP(const GPRegister& rd, const int32_t& imm32) { immediateGroup1(rd, imm32, 0x07); }

// ----------------------------------------------

void x64Emitter::PUSH(const GPRegister& rs)
{
    if (rs > RDI)
        *m_pbWrite++ = 0x41; // REX

    *m_pbWrite++ = 0x50 + (rs.id & 0x07);
}

void x64Emitter::POP(const GPRegister& rd)
{
    if (rd > RDI)
        *m_pbWrite++ = 0x41; // REX

    *m_pbWrite++ = 0x58 + (rd.id & 0x07);
}

// ----------------------------------------------

void x64Emitter::LABEL(Label* label)
{
    *label = m_pbWrite;
}

void x64Emitter::placeJump(const JMP_Handle& jump, const Label& label)
{
    intptr_t nOffset = label - reinterpret_cast<uint8_t*>(jump);

    *jump = nOffset - 4;
}

// ----------------------------------------------

void x64Emitter::jumpBack(const Label& label, uint8_t op8, uint16_t op32)
{
    int64_t nOffset = label - m_pbWrite;

    // -128 <= X - 2 <= 127
    if (nOffset >= -126 && nOffset <= 129)
    {
        *m_pbWrite++ = op8;
        *m_pbWrite++ = (nOffset - 2) & 0xFF;
    }
    else
    {
        *m_pbWrite++ = op32 & 0xFF;
        *m_pbWrite++ = (op32 >> 8) & 0xFF;
        
        *(reinterpret_cast<int32_t*>(m_pbWrite)) = (nOffset - 6) & 0xFFFFFFFF;
        m_pbWrite += 4;
    }
}

JMP_Handle x64Emitter::jumpFuture(uint16_t op32)
{
    *m_pbWrite++ = op32 & 0xFF;
    *m_pbWrite++ = (op32 >> 8) & 0xFF;

    JMP_Handle offset = reinterpret_cast<JMP_Handle>(m_pbWrite);
    m_pbWrite += 4;

    return offset;
}

// ----------------------------------------------

#define JMP_INSTRUCTION(name, op8, op32) \
    void x64Emitter::name(const Label& label) { jumpBack(label, op8, op32); } \
    JMP_Handle x64Emitter::name() { return jumpFuture(op32); }


JMP_INSTRUCTION(JA , 0x77, 0x870F)
JMP_INSTRUCTION(JAE, 0x73, 0x830F)
JMP_INSTRUCTION(JB , 0x72, 0x820F)
JMP_INSTRUCTION(JBE, 0x76, 0x860F)
JMP_INSTRUCTION(JC , 0x72, 0x820F)
JMP_INSTRUCTION(JE , 0x74, 0x840F)
JMP_INSTRUCTION(JG , 0x7F, 0x8F0F)
JMP_INSTRUCTION(JGE, 0x7D, 0x8D0F)
JMP_INSTRUCTION(JL , 0x7C, 0x8C0F)
JMP_INSTRUCTION(JLE, 0x7E, 0x8E0F)
JMP_INSTRUCTION(JNA, 0x76, 0x860F)
JMP_INSTRUCTION(JNAE, 0x72, 0x820F)
JMP_INSTRUCTION(JNB, 0x73, 0x830F)
JMP_INSTRUCTION(JNBE, 0x77, 0x870F)
JMP_INSTRUCTION(JNC, 0x73, 0x830F)
JMP_INSTRUCTION(JNE, 0x75, 0x850F)
JMP_INSTRUCTION(JNG, 0x7E, 0x8E0F)
JMP_INSTRUCTION(JNGE, 0x7C, 0x8C0F)
JMP_INSTRUCTION(JNL, 0x7D, 0x8D0F)
JMP_INSTRUCTION(JNLE, 0x7F, 0x8F0F)
JMP_INSTRUCTION(JNO, 0x71, 0x810F)
JMP_INSTRUCTION(JNP, 0x7B, 0x8B0F)
JMP_INSTRUCTION(JNS, 0x79, 0x890F)
JMP_INSTRUCTION(JNZ, 0x75, 0x850F)
JMP_INSTRUCTION(JO , 0x70, 0x800F)
JMP_INSTRUCTION(JP , 0x7A, 0x8A0F)
JMP_INSTRUCTION(JPE, 0x7A, 0x8A0F)
JMP_INSTRUCTION(JPO, 0x7B, 0x8B0F)
JMP_INSTRUCTION(JS , 0x78, 0x880F)
JMP_INSTRUCTION(JZ , 0x74, 0x840F)

// ----------------------------------------------

void x64Emitter::JMP(const Label& label) 
{ 
    int64_t nOffset = label - m_pbWrite;

    // -128 <= X - 2 <= 127
    if (nOffset >= -126 && nOffset <= 129)
    {
        *m_pbWrite++ = 0xEB;
        *m_pbWrite++ = (nOffset - 2) & 0xFF;
    }
    else
    {
        *m_pbWrite++ = 0xE9;
        
        *(reinterpret_cast<int32_t*>(m_pbWrite)) = (nOffset - 5) & 0xFFFFFFFF;
        m_pbWrite += 4;
    }
}

JMP_Handle x64Emitter::JMP() 
{ 
    *m_pbWrite++ = 0xE9;

    JMP_Handle offset = reinterpret_cast<JMP_Handle>(m_pbWrite);
    m_pbWrite += 4;

    return offset;
}

// ----------------------------------------------

void x64Emitter::CALL(const Label& label) 
{ 
    int64_t nOffset = label - m_pbWrite;

    *m_pbWrite++ = 0xE8;
        
    *(reinterpret_cast<int32_t*>(m_pbWrite)) = (nOffset - 5) & 0xFFFFFFFF;
    m_pbWrite += 4;
}

JMP_Handle x64Emitter::CALL() 
{ 
    *m_pbWrite++ = 0xE8;

    JMP_Handle offset = reinterpret_cast<JMP_Handle>(m_pbWrite);
    m_pbWrite += 4;

    return offset;
}



